package app

import (
	"encoding/csv"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/cio"
	"github.com/containerd/containerd/containers"
	"github.com/containerd/containerd/contrib/apparmor"
	"github.com/containerd/containerd/contrib/nvidia"
	"github.com/containerd/containerd/contrib/seccomp"
	"github.com/containerd/containerd/namespaces"
	"github.com/containerd/containerd/oci"
	"github.com/containerd/containerd/runtime/v2/runc/options"
	"github.com/containerd/containerd/snapshots"
	"github.com/opencontainers/runtime-spec/specs-go"
	"github.com/urfave/cli/v2"

	gocontext "context"
)

func AppContext(context *cli.Context) (gocontext.Context, gocontext.CancelFunc) {
	ctx := gocontext.Background()
	// grpc header add namespace
	ctx = namespaces.WithNamespace(ctx, context.String("namespace"))

	var cancel gocontext.CancelFunc
	ctx, cancel = gocontext.WithCancel(ctx)

	return ctx, cancel
}

func NewClient(context *cli.Context) (*containerd.Client, gocontext.Context, gocontext.CancelFunc, error) {
	client, err := containerd.New(context.String("address"))
	if err != nil {
		return nil, nil, nil, err
	}
	ctx, cancel := AppContext(context)
	return client, ctx, cancel, nil
}

func NewContainer(ctx gocontext.Context, client *containerd.Client, context *cli.Context) (containerd.Container, error) {
	ref := context.Args().First()
	id := context.Args().Get(1)
	args := context.Args().Slice()[2:]

	var (
		opts  []oci.SpecOpts
		cOpts []containerd.NewContainerOpts
		spec  containerd.NewContainerOpts
	)

	opts = append(opts, oci.WithDefaultSpec(), oci.WithDefaultUnixDevices)
	opts = append(opts, oci.WithEnv(context.StringSlice("env")))
	opts = append(opts, withMounts(context))

	if context.Bool("rootfs") {
		rootfs, err := filepath.Abs(ref)
		if err != nil {
			return nil, err
		}
		opts = append(opts, oci.WithRootFSPath(rootfs))
		cOpts = append(cOpts, containerd.WithContainerLabels(LabelArgs(context.StringSlice("label"))))
	} else {
		snapshotter := context.String("snapshotter")
		var image containerd.Image
		i, err := client.ImageService().Get(ctx, ref)
		if err != nil {
			// image not found
			image, err = client.Pull(ctx, ref)
			if err != nil {
				return nil, err
			}
		} else {
			image = containerd.NewImage(client, i)
		}

		unpacked, err := image.IsUnpacked(ctx, snapshotter)
		if err != nil {
			return nil, err
		}
		if !unpacked {
			if err := image.Unpack(ctx, snapshotter); err != nil {
				return nil, err
			}
		}
		labels := buildLabels(LabelArgs(context.StringSlice("label")), image.Labels())
		opts = append(opts, oci.WithImageConfig(image))
		cOpts = append(cOpts,
			containerd.WithImage(image),
			containerd.WithImageConfigLabels(image),
			containerd.WithAdditionalContainerLabels(labels),
			containerd.WithSnapshotter(snapshotter))
		if uidmap, gidmap := context.String("uidmap"), context.String("gidmap"); uidmap != "" && gidmap != "" {
			uidMap, err := parseIDMapping(uidmap)
			if err != nil {
				return nil, err
			}
			gidMap, err := parseIDMapping(gidmap)
			if err != nil {
				return nil, err
			}
			opts = append(opts,
				oci.WithUserNamespace([]specs.LinuxIDMapping{uidMap}, []specs.LinuxIDMapping{gidMap}))
			if context.Bool("remap-labels") {
				cOpts = append(cOpts, containerd.WithNewSnapshot(id, image,
					containerd.WithRemapperLabels(0, uidMap.HostID, 0, gidMap.HostID, uidMap.Size)))
			} else {
				cOpts = append(cOpts, containerd.WithRemappedSnapshot(id, image, uidMap.HostID, gidMap.HostID))
			}
		} else {
			cOpts = append(cOpts, containerd.WithNewSnapshot(id, image,
				snapshots.WithLabels(LabelArgs(context.StringSlice("snapshotter-label")))))
		}
		cOpts = append(cOpts, containerd.WithImageStopSignal(image, "SIGTERM"))
	}
	if len(args) > 0 {
		opts = append(opts, oci.WithProcessArgs(args...))
	}
	if cwd := context.String("cwd"); cwd != "" {
		opts = append(opts, oci.WithProcessCwd(cwd))
	}
	if user := context.String("user"); user != "" {
		opts = append(opts, oci.WithUser(user), oci.WithAdditionalGIDs(user))
	}

	privileged := context.Bool("privileged")
	privilegedWithoutHostDevices := context.Bool("privileged-without-host-devices")
	if privilegedWithoutHostDevices && !privileged {
		return nil, errors.New("can't use 'privileged-without-host-devices' without 'privileged' specified")
	}
	if privileged {
		if privilegedWithoutHostDevices {
			opts = append(opts, oci.WithPrivileged)
		} else {
			opts = append(opts, oci.WithPrivileged, oci.WithAllDevicesAllowed, oci.WithHostDevices)
		}
	}

	if context.Bool("net-host") {
		hostname, err := os.Hostname()
		if err != nil {
			return nil, fmt.Errorf("get hostname: %w", err)
		}
		opts = append(opts,
			oci.WithHostNamespace(specs.NetworkNamespace),
			oci.WithHostHostsFile,
			oci.WithHostResolvconf,
			oci.WithEnv([]string{fmt.Sprintf("HOSTNAME=%s", hostname)}),
		)
	}
	if annoStrings := context.StringSlice("annotation"); len(annoStrings) > 0 {
		annos, err := AnnotationArgs(annoStrings)
		if err != nil {
			return nil, err
		}
		opts = append(opts, oci.WithAnnotations(annos))
	}

	if caps := context.StringSlice("cap-add"); len(caps) > 0 {
		for _, cap := range caps {
			if !strings.HasPrefix(cap, "CAP_") {
				return nil, errors.New("capabilities must be specified with 'CAP_' prefix")
			}
		}
		opts = append(opts, oci.WithAddedCapabilities(caps))
	}

	if caps := context.StringSlice("cap-drop"); len(caps) > 0 {
		for _, cap := range caps {
			if !strings.HasPrefix(cap, "CAP_") {
				return nil, errors.New("capabilities must be specified with 'CAP_' prefix")
			}
		}
		opts = append(opts, oci.WithDroppedCapabilities(caps))
	}

	seccompProfile := context.String("seccomp-profile")
	if !context.Bool("seccomp") && seccompProfile != "" {
		return nil, errors.New("seccomp must be set to true, if using a custom seccomp-profile")
	}

	if context.Bool("seccomp") {
		if seccompProfile != "" {
			opts = append(opts, seccomp.WithProfile(seccompProfile))
		} else {
			opts = append(opts, seccomp.WithDefaultProfile())
		}
	}

	if s := context.String("apparmor-default-profile"); len(s) > 0 {
		opts = append(opts, apparmor.WithDefaultProfile(s))
	}

	if s := context.String("apparmor-profile"); len(s) > 0 {
		if len(context.String("apparmor-default-profile")) > 0 {
			return nil, errors.New("apparmor-profile conflicts with apparmor-default-profile")
		}
		opts = append(opts, apparmor.WithProfile(s))
	}

	if cpus := context.Float64("cpus"); cpus > 0.0 {
		var (
			period = uint64(100000)
			quota  = int64(cpus * 100000.0)
		)
		opts = append(opts, oci.WithCPUCFS(quota, period))
	}

	if cpusetCpus := context.String("cpuset-cpus"); len(cpusetCpus) > 0 {
		opts = append(opts, oci.WithCPUs(cpusetCpus))
	}

	if cpusetMems := context.String("cpuset-mems"); len(cpusetMems) > 0 {
		opts = append(opts, oci.WithCPUsMems(cpusetMems))
	}

	if shares := context.Int("cpu-shares"); shares > 0 {
		opts = append(opts, oci.WithCPUShares(uint64(shares)))
	}

	quota := context.Int64("cpu-quota")
	period := context.Uint64("cpu-period")
	if quota != -1 || period != 0 {
		if cpus := context.Float64("cpus"); cpus > 0.0 {
			return nil, errors.New("cpus and quota/period should be used separately")
		}
		opts = append(opts, oci.WithCPUCFS(quota, period))
	}

	joinNs := context.StringSlice("with-ns")
	for _, ns := range joinNs {
		nsType, nsPath, ok := strings.Cut(ns, ":")
		if !ok {
			return nil, errors.New("joining a Linux namespace using --with-ns requires the format 'nstype:path'")
		}
		if !validNamespace(nsType) {
			return nil, errors.New("the Linux namespace type specified in --with-ns is not valid: " + nsType)
		}
		opts = append(opts, oci.WithLinuxNamespace(specs.LinuxNamespace{
			Type: specs.LinuxNamespaceType(nsType),
			Path: nsPath,
		}))
	}
	if context.IsSet("gpus") {
		opts = append(opts, nvidia.WithGPUs(nvidia.WithDevices(context.IntSlice("gpus")...), nvidia.WithAllCapabilities))
	}
	if context.IsSet("allow-new-privs") {
		opts = append(opts, oci.WithNewPrivileges)
	}
	if context.IsSet("cgroup") {
		// NOTE: can be set to "" explicitly for disabling cgroup.
		opts = append(opts, oci.WithCgroup(context.String("cgroup")))
	}

	limit := context.Uint64("memory-limit")
	if limit != 0 {
		opts = append(opts, oci.WithMemoryLimit(limit))
	}

	rootfsPropagation := context.String("rootfs-propagation")
	if rootfsPropagation != "" {
		opts = append(opts, func(_ gocontext.Context, _ oci.Client, _ *containers.Container, s *oci.Spec) error {
			if s.Linux != nil {
				s.Linux.RootfsPropagation = rootfsPropagation
			} else {
				s.Linux = &specs.Linux{
					RootfsPropagation: rootfsPropagation,
				}
			}

			return nil
		})
	}

	if c := context.String("rdt-class"); c != "" {
		opts = append(opts, oci.WithRdt(c, "", ""))
	}
	if hostname := context.String("hostname"); hostname != "" {
		opts = append(opts, oci.WithHostname(hostname))
	}

	cOpts = append(cOpts, containerd.WithRuntime("io.containerd.runc.v2", &options.Options{}))
	opts = append(opts, oci.WithAnnotations(LabelArgs(context.StringSlice("label"))))
	var s specs.Spec
	spec = containerd.WithSpec(&s, opts...)

	cOpts = append(cOpts, spec)

	return client.NewContainer(ctx, id, cOpts...)
}

func NewTask(ctx gocontext.Context, client *containerd.Client, container containerd.Container, nullIO bool, ioOpts []cio.Opt, opts ...containerd.NewTaskOpts) (containerd.Task, error) {
	stdinC := &stdinCloser{
		stdin: os.Stdin,
	}
	spec, err := container.Spec(ctx)
	if err != nil {
		return nil, err
	}
	if spec.Linux != nil {
		if len(spec.Linux.UIDMappings) != 0 {
			opts = append(opts, containerd.WithUIDOwner(spec.Linux.UIDMappings[0].HostID))
		}
		if len(spec.Linux.GIDMappings) != 0 {
			opts = append(opts, containerd.WithGIDOwner(spec.Linux.GIDMappings[0].HostID))
		}
	}

	var ioCreator cio.Creator
	if nullIO {
		ioCreator = cio.NullIO
	} else {
		ioCreator = cio.NewCreator(append([]cio.Opt{cio.WithStreams(stdinC, os.Stdout, os.Stderr)}, ioOpts...)...)
	}

	t, err := container.NewTask(ctx, ioCreator, opts...)
	if err != nil {
		return nil, err
	}
	stdinC.closer = func() {
		t.CloseIO(ctx, containerd.WithStdinCloser)
	}
	return t, nil
}

func withMounts(context *cli.Context) oci.SpecOpts {
	return func(ctx gocontext.Context, client oci.Client, container *containers.Container, s *specs.Spec) error {
		mounts := make([]specs.Mount, 0)
		dests := make([]string, 0)
		for _, mount := range context.StringSlice("mount") {
			m, err := parseMountFlag(mount)
			if err != nil {
				return err
			}
			mounts = append(mounts, m)
			dests = append(dests, m.Destination)
		}
		return oci.Compose(oci.WithoutMounts(dests...), oci.WithMounts(mounts))(ctx, client, container, s)
	}
}

func parseMountFlag(m string) (specs.Mount, error) {
	mount := specs.Mount{}
	r := csv.NewReader(strings.NewReader(m))

	fields, err := r.Read()
	if err != nil {
		return mount, err
	}

	for _, field := range fields {
		key, val, ok := strings.Cut(field, "=")
		if !ok {
			return mount, errors.New("invalid mount specification: expected key=val")
		}

		switch key {
		case "type":
			mount.Type = val
		case "source", "src":
			mount.Source = val
		case "destination", "dst":
			mount.Destination = val
		case "options":
			mount.Options = strings.Split(val, ":")
		default:
			return mount, fmt.Errorf("mount option %q not supported", key)
		}
	}

	return mount, nil
}
